import { useEffect, useRef, useState } from "react";
import { SocketIOContext } from "./context";
import { useContextSelector } from 'use-context-selector';
import { Socket } from "socket.io-client";

interface UseSocketReturn {
  socket: Socket | undefined;
  unsubscribe: () => void;
  subscribe: () => void
}

export const useSocket = (
  eventKey?: string,
  callback?: (...args: any) => void
): UseSocketReturn => {
  const { socket } = useContextSelector(SocketIOContext, (s)=>s);
  const callbackRef = useRef(callback);

  callbackRef.current = callback;

  const socketHandlerRef = useRef((...args: any[]) => {
    if (callbackRef.current) {
      callbackRef.current(...args);
    }
  });

  const subscribe = () => {
    if (eventKey) {
      socket?.on(eventKey, socketHandlerRef.current);
    }
  };

  const unsubscribe = () => {
    if (eventKey) {
      socket?.removeListener(eventKey, socketHandlerRef.current);
    }
  };

  useEffect(() => {
    subscribe();

    return unsubscribe;
  }, [eventKey]);

  return { socket, unsubscribe, subscribe };
};

interface LastMessage {
  data: string;
  socket: any; // Replace 'any' with the actual type of your socket if available
  unsubscribe: () => void;
  subscribe: () => void;
}

export const useLastMessage = (eventKey: string) : LastMessage => {
  const { socket } = useContextSelector(SocketIOContext, (s)=>s);
  const [data, setData] = useState("");

  const subscribe = () => {
    if (eventKey) {
      socket?.on(eventKey, setData);
    }
  };

  const unsubscribe = () => {
    if (eventKey) {
      socket?.removeListener(eventKey, setData);
    }
  };

  useEffect(() => {
    subscribe();

    return unsubscribe;
  }, [eventKey]);

  return { data, socket, unsubscribe, subscribe };
};
